use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{Mutex, mpsc};
use tokio::net::{TcpListener, TcpStream};
use tokio_tungstenite::{accept_async, tungstenite::Message};
use futures_util::{SinkExt, StreamExt};
use serde_json::{json, Value};
use log::{info, error, debug};



/// WebSocket连接信息
#[derive(Debug, Clone)]
pub struct Connection {
    #[allow(dead_code)]
    pub id: String,
    pub device_type: Option<String>, // "tv" 或 "mobile"
    pub room_id: Option<String>,
    pub device_id: Option<String>,
    pub device_name: Option<String>, // 设备名称
    pub sender: Option<mpsc::UnboundedSender<Message>>,
}

/// WebSocket服务器
pub struct WebSocketServer {
    connections: Arc<Mutex<HashMap<String, Connection>>>,
    rooms: Arc<Mutex<HashMap<String, Room>>>,
}

/// 游戏房间
#[derive(Debug, Clone)]
struct Room {
    #[allow(dead_code)]
    id: String,
    tv_connection: Option<String>,
    mobile_connections: Vec<String>,
}

impl WebSocketServer {
    pub fn new() -> Self {
        Self {
            connections: Arc::new(Mutex::new(HashMap::new())),
            rooms: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// 启动WebSocket服务器
    pub async fn start(&self, port: u16) -> Result<(), String> {
        let addr = format!("0.0.0.0:{}", port);
        let listener = TcpListener::bind(&addr).await.map_err(|e| e.to_string())?;
        
        info!("WebSocket服务器已启动，监听端口: {}", port);
        
        while let Ok((stream, _)) = listener.accept().await {
            let connections = Arc::clone(&self.connections);
            let rooms = Arc::clone(&self.rooms);
            
            tokio::spawn(async move {
                let _ = handle_connection(stream, connections, rooms).await;
            });
        }

        Ok(())
    }
}

/// 处理单个WebSocket连接
async fn handle_connection(
    stream: TcpStream,
    connections: Arc<Mutex<HashMap<String, Connection>>>,
    rooms: Arc<Mutex<HashMap<String, Room>>>,
) {
    let ws_stream = match accept_async(stream).await {
        Ok(ws) => ws,
        Err(_) => return,
    };
    let (mut ws_sender, mut ws_receiver) = ws_stream.split();
    
    let connection_id = generate_id();
    info!("新连接: {}", connection_id);

    // 创建消息发送通道
    let (tx, mut rx) = mpsc::unbounded_channel::<Message>();

    // 添加连接
    {
        let mut conns = connections.lock().await;
        conns.insert(connection_id.clone(), Connection {
            id: connection_id.clone(),
            device_type: None,
            room_id: None,
            device_id: None,
            device_name: None,
            sender: Some(tx),
        });
    }

    // 发送连接确认
    let welcome_msg = json!({
        "type": "connection_established",
        "connectionId": connection_id
    });
    
    if let Err(e) = ws_sender.send(Message::Text(welcome_msg.to_string())).await {
        error!("发送欢迎消息失败: {}", e);
        return;
    }

    // 启动消息发送任务
    let connection_id_clone = connection_id.clone();
    tokio::spawn(async move {
        while let Some(message) = rx.recv().await {
            if let Err(e) = ws_sender.send(message).await {
                error!("发送消息失败 to {}: {}", connection_id_clone, e);
                break;
            }
        }
    });

    // 处理消息
    while let Some(msg) = ws_receiver.next().await {
        match msg {
            Ok(Message::Text(text)) => {
                // 解析消息
                let msg: Value = match serde_json::from_str(&text) {
                    Ok(msg) => msg,
                    Err(_) => continue, // 简单跳过无效消息
                };

                let msg_type = msg["type"].as_str().unwrap_or("unknown");
                
                // 只记录非绘图数据的消息，减少日志输出
                if msg_type != "draw_data" {
                    info!("收到消息类型: {} from {}", msg_type, connection_id);
                }
                
                if let Err(e) = handle_message(
                    &connection_id,
                    msg,
                    &connections,
                    &rooms,
                ).await {
                    error!("处理消息失败: {}", e);
                }
            }
            Ok(Message::Close(_)) => {
                info!("连接 {} 主动关闭", connection_id);
                break;
            }
            Err(e) => {
                error!("WebSocket错误: {}", e);
                break;
            }
            _ => {}
        }
    }

    // 清理连接
    cleanup_connection(&connection_id, &connections, &rooms).await;
    info!("连接 {} 已断开", connection_id);
}

/// 处理WebSocket消息
async fn handle_message(
    connection_id: &str,
    data: Value,
    connections: &Arc<Mutex<HashMap<String, Connection>>>,
    rooms: &Arc<Mutex<HashMap<String, Room>>>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let msg_type = data["type"].as_str().unwrap_or("unknown");
    
    match msg_type {
        "register_tv" => {
            handle_register_tv(connection_id, data, connections, rooms).await?
        }
        "register_mobile" => {
            handle_register_mobile(connection_id, data, connections, rooms).await?
        }
        "draw_data" => {
            handle_draw_data(connection_id, data, connections, rooms).await?
        }
        "guess_submitted" => {
            handle_guess_submitted(connection_id, data, connections, rooms).await?
        }
        "clear_canvas" => {
            handle_clear_canvas(connection_id, data, connections, rooms).await?
        }
        "start_game" => {
            handle_start_game(connection_id, data, connections, rooms).await?
        }
        _ => {
            // 只记录真正未知的消息类型
            if msg_type != "draw_data" {
                debug!("未知消息类型: {}", msg_type);
            }
        }
    }
    
    Ok(())
}

/// 发送消息到指定连接
async fn send_message_to_connection(
    connection_id: &str,
    message: Message,
    connections: &Arc<Mutex<HashMap<String, Connection>>>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let conns = connections.lock().await;
    if let Some(conn) = conns.get(connection_id) {
        if let Some(sender) = &conn.sender {
            sender.send(message).map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
        }
    }
    Ok(())
}

/// 处理电视端注册
async fn handle_register_tv(
    connection_id: &str,
    _data: Value,
    connections: &Arc<Mutex<HashMap<String, Connection>>>,
    rooms: &Arc<Mutex<HashMap<String, Room>>>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let room_id = generate_id();
    
    // 更新连接信息
    {
        let mut conns = connections.lock().await;
        if let Some(conn) = conns.get_mut(connection_id) {
            conn.device_type = Some("tv".to_string());
            conn.room_id = Some(room_id.clone());
        }
    }
    
    // 创建房间
    {
        let mut room_map = rooms.lock().await;
        room_map.insert(room_id.clone(), Room {
            id: room_id.clone(),
            tv_connection: Some(connection_id.to_string()),
            mobile_connections: Vec::new(),
        });
    }
    
    info!("电视端注册: {}, 房间: {}", connection_id, room_id);
    
    // 发送注册成功消息
    let response = json!({
        "type": "tv_registered",
        "roomId": room_id
    });
    
    send_message_to_connection(connection_id, Message::Text(response.to_string()), connections).await?;
    Ok(())
}

/// 处理手机端注册
async fn handle_register_mobile(
    connection_id: &str,
    data: Value,
    connections: &Arc<Mutex<HashMap<String, Connection>>>,
    rooms: &Arc<Mutex<HashMap<String, Room>>>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let room_id = data["roomId"].as_str().unwrap_or("").to_string();
    let device_id = data["deviceId"].as_str().unwrap_or(&generate_id()).to_string();
    let device_name = data["deviceName"].as_str().unwrap_or("未知设备").to_string();
    
    // 检查房间是否存在
    let tv_connection = {
        let room_map = rooms.lock().await;
        if let Some(room) = room_map.get(&room_id) {
            room.tv_connection.clone()
        } else {
            let error_msg = json!({
                "type": "error",
                "message": "Room not found"
            });
            send_message_to_connection(connection_id, Message::Text(error_msg.to_string()), connections).await?;
            return Ok(());
        }
    };
    
    // 更新连接信息
    {
        let mut conns = connections.lock().await;
        if let Some(conn) = conns.get_mut(connection_id) {
            conn.device_type = Some("mobile".to_string());
            conn.room_id = Some(room_id.clone());
            conn.device_id = Some(device_id.clone());
            conn.device_name = Some(device_name.clone());
        }
    }
    
    // 添加到房间
    {
        let mut room_map = rooms.lock().await;
        if let Some(room) = room_map.get_mut(&room_id) {
            room.mobile_connections.push(connection_id.to_string());
        }
    }
    
    info!("手机端注册: {}, 设备ID: {}, 房间: {}", connection_id, device_id, room_id);
    
    // 通知电视端有新设备连接
    if let Some(tv_conn_id) = tv_connection {
        let tv_msg = json!({
            "type": "device_connected",
            "deviceId": device_id,
            "deviceName": device_name,
            "connectionId": connection_id
        });
        
        // 发送消息到电视端
        if let Err(e) = send_message_to_connection(&tv_conn_id, Message::Text(tv_msg.to_string()), connections).await {
            error!("发送device_connected消息到TV端失败: {}", e);
        } else {
            info!("已通知TV端设备连接: {}", device_id);
        }
    }
    
    // 发送注册成功消息
    let response = json!({
        "type": "mobile_registered",
        "deviceId": device_id,
        "roomId": room_id
    });
    
    send_message_to_connection(connection_id, Message::Text(response.to_string()), connections).await?;
    Ok(())
}

/// 处理绘画数据
async fn handle_draw_data(
    connection_id: &str,
    data: Value,
    connections: &Arc<Mutex<HashMap<String, Connection>>>,
    rooms: &Arc<Mutex<HashMap<String, Room>>>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    // 获取连接信息
    let (room_id, device_id) = {
        let conns = connections.lock().await;
        if let Some(conn) = conns.get(connection_id) {
            (conn.room_id.clone(), conn.device_id.clone())
        } else {
            return Ok(());
        }
    };
    
    if let (Some(room_id), Some(device_id)) = (room_id, device_id) {
        // 获取TV端连接ID
        let tv_connection_id = {
            let room_map = rooms.lock().await;
            if let Some(room) = room_map.get(&room_id) {
                room.tv_connection.clone()
            } else {
                debug!("房间 {} 不存在", room_id);
                return Ok(());
            }
        };
        
        if let Some(tv_conn_id) = tv_connection_id {
            // 转发绘画数据到电视端，保留所有原始字段
            let forward_msg = json!({
                "type": "draw_data",
                "deviceId": device_id,
                "drawType": data["drawType"],
                "lastX": data["lastX"],
                "lastY": data["lastY"],
                "x": data["x"],
                "y": data["y"],
                "color": data["color"],
                "size": data["size"],
                "isDrawing": data["isDrawing"],
                "isEraser": data["isEraser"]
            });
            
            // 发送消息到TV端
            if let Err(e) = send_message_to_connection(&tv_conn_id, Message::Text(forward_msg.to_string()), connections).await {
                error!("发送绘画数据到TV端失败: {}", e);
            }
        } else {
            debug!("房间 {} 没有TV端连接", room_id);
        }
    }
    
    Ok(())
}

/// 处理猜测提交
async fn handle_guess_submitted(
    connection_id: &str,
    data: Value,
    connections: &Arc<Mutex<HashMap<String, Connection>>>,
    _rooms: &Arc<Mutex<HashMap<String, Room>>>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let (room_id, device_id) = {
        let conns = connections.lock().await;
        if let Some(conn) = conns.get(connection_id) {
            (conn.room_id.clone(), conn.device_id.clone())
        } else {
            return Ok(());
        }
    };
    
    if let (Some(room_id), Some(device_id)) = (room_id, device_id) {
        let forward_msg = json!({
            "type": "guess_submitted",
            "deviceId": device_id,
            "guess": data["guess"]
        });
        
        debug!("转发猜测到房间 {}: {:?}", room_id, forward_msg);
        // 在实际实现中，需要发送到电视端
    }
    
    Ok(())
}

/// 处理清除画布
async fn handle_clear_canvas(
    connection_id: &str,
    _data: Value,
    _connections: &Arc<Mutex<HashMap<String, Connection>>>,
    _rooms: &Arc<Mutex<HashMap<String, Room>>>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    debug!("处理清除画布请求 from {}", connection_id);
    // 实现清除画布逻辑
    Ok(())
}

/// 处理游戏开始
async fn handle_start_game(
    connection_id: &str,
    data: Value,
    _connections: &Arc<Mutex<HashMap<String, Connection>>>,
    _rooms: &Arc<Mutex<HashMap<String, Room>>>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    debug!("处理游戏开始请求 from {}: {:?}", connection_id, data);
    // 实现游戏开始逻辑
    Ok(())
}

/// 清理连接
async fn cleanup_connection(
    connection_id: &str,
    connections: &Arc<Mutex<HashMap<String, Connection>>>,
    rooms: &Arc<Mutex<HashMap<String, Room>>>,
) {
    let connection = {
        let mut conns = connections.lock().await;
        conns.remove(connection_id)
    };
    
    if let Some(conn) = connection {
        if let Some(room_id) = conn.room_id {
            let mut room_map = rooms.lock().await;
            if let Some(room) = room_map.get_mut(&room_id) {
                if conn.device_type.as_deref() == Some("tv") {
                    // 电视端断开，删除房间
                    room_map.remove(&room_id);
                    info!("房间 {} 已删除", room_id);
                } else if conn.device_type.as_deref() == Some("mobile") {
                    // 手机端断开，从房间中移除
                    room.mobile_connections.retain(|id| id != connection_id);
                    info!("手机端 {} 从房间 {} 中移除", connection_id, room_id);
                }
            }
        }
    }
}

/// 生成唯一ID
fn generate_id() -> String {
    use std::time::{SystemTime, UNIX_EPOCH};
    let timestamp = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_millis();
    format!("{:x}{:x}", timestamp, rand::random::<u32>())
}

/// 启动WebSocket服务器的公共函数
pub async fn start_websocket_server(port: u16) -> Result<(), String> {
    let server = WebSocketServer::new();
    server.start(port).await
}
